# File: fedentgate/core/entropy_gating.py
"""Entropy-aware Gating Module for Client Selection"""

import torch
import numpy as np

class EntropyGating:
    def __init__(self, base_threshold=1.2, decay_factor=0.95):
        """
        Initialize entropy-based gating mechanism
        Args:
            base_threshold: Initial entropy threshold (τ_base)
            decay_factor: Threshold decay coefficient (γ)
        """
        self.base_threshold = base_threshold
        self.decay_factor = decay_factor
        self.current_thresholds = {}

    def compute_entropy(self, data_distribution):
        """
        Calculate information entropy for client data
        Args:
            data_distribution: Tensor of class probabilities [C]
        Returns:
            entropy: Scalar entropy value (H_k)
        """
        return -torch.sum(data_distribution * torch.log(data_distribution + 1e-9))

    def update_threshold(self, client_id, current_round, total_rounds):
        """
        Update dynamic entropy threshold for client
        Args:
            client_id: Client identifier
            current_round: Current training round (t)
            total_rounds: Total training rounds (T)
        """
        self.current_thresholds[client_id] = self.base_threshold + \
            (current_round / total_rounds) * self.decay_factor

    def gating_decision(self, entropy, client_id):
        """
        Make binary gating decision based on entropy
        Args:
            entropy: Computed entropy value (H_k)
            client_id: Client identifier
        Returns:
            gating_state: Binary decision (0/1)
        """
        threshold = self.current_thresholds.get(client_id, self.base_threshold)
        return 1 if entropy > threshold else 0
